# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import functools
from types import LambdaType

from coin.exchange.base.order_gateway import OrderSide
from coin.strategy.mm.simple_sim.base import (PositionCalculator,
                                              SimpleMatcher,
                                              UnavailabilityChecker)


class PassiveSimExecutor(object):
  def __init__(self,
               lot_size,
               min_pos,
               max_pos,
               maker_fee,
               taker_fee,
               execution_delay,
               post_only,
               *,
               min_fire_period=2 * (10**9),
               trade_qty_func=None,
               ignore_book_fill=True,
               do_agg=False,
               unavail_checker=None,
               unavail_intervals_file=None):
    self._get_lot_size = None
    self._get_min_pos = None
    self._get_max_pos = None

    self._last_sell_timestamp = 0
    self._last_buy_timestamp = 0

    if isinstance(lot_size, LambdaType):
      self._get_lot_size = lot_size
    else:
      self._get_lot_size = lambda x: lot_size
    if isinstance(min_pos, LambdaType):
      self._get_min_pos = min_pos
    else:
      self._get_min_pos = lambda x: min_pos
    if isinstance(max_pos, LambdaType):
      self._get_max_pos = max_pos
    else:
      self._get_max_pos = lambda x: max_pos

    self._min_fire_period = min_fire_period
    self._trade_qty_func = trade_qty_func or (lambda trade: trade.qty)

    self.fill_callback = lambda price, qty, side, maker, timestamp: None
    self.pos_callback = lambda pos, price, qty, side, maker, timestamp: None

    self.pos = PositionCalculator(maker_fee, taker_fee)

    unavail_checker = (unavail_checker
                       or UnavailabilityChecker.get_from_file(unavail_intervals_file))
    self._ask_matcher = SimpleMatcher(OrderSide.SELL,
                                      execution_delay,
                                      post_only,
                                      ignore_book_fill=ignore_book_fill,
                                      unavail_checker=unavail_checker)
    self._bid_matcher = SimpleMatcher(OrderSide.BUY,
                                      execution_delay,
                                      post_only,
                                      ignore_book_fill=ignore_book_fill,
                                      unavail_checker=unavail_checker)
    self._last_fire_time = 0

    self._max_agg_buy_price = None
    self._min_agg_sell_price = None

    if do_agg:
      # TODO(inkyu): Let's implement this later
      assert not isinstance(lot_size, LambdaType)

      self._agg_executor = AggressionSimExecutor(maker_fee,
                                                 taker_fee,
                                                 execution_delay,
                                                 lot_size,
                                                 True,
                                                 trade_qty_func,
                                                 unavail_checker=unavail_checker)
      self._agg_executor.fill_callback = self._on_agg_fill
      self._last_agg_ts = 0
    else:
      self._agg_executor = None

  '''Aggression'''

  def agg_buy(self, price, timestamp):
    assert self._agg_executor is not None

    if (self.pos.pos < self._get_max_pos(timestamp)
        and timestamp - self._last_agg_ts >= self._min_fire_period):
      self._agg_executor.buy(timestamp, price, self._get_lot_size(timestamp))
      self._last_agg_ts = timestamp

  def agg_sell(self, price, timestamp):
    assert self._agg_executor is not None

    if (self.pos.pos > self._get_min_pos(timestamp)
        and timestamp - self._last_agg_ts >= self._min_fire_period):
      self._agg_executor.sell(timestamp, price, self._get_lot_size(timestamp))
      self._last_agg_ts = timestamp

  def cancel_agg_orders(self, timestamp):
    self._agg_executor.cancel_all(timestamp)

  def _on_agg_fill(self, price, qty, side, maker, timestamp):
    self.fill_callback(price, qty, side, maker, timestamp)
    self.pos.update(price, qty, side, maker)
    self.pos_callback(self.pos, price, qty, side, maker, timestamp)

  def _within_sell_update_period(self, timestamp):
    return (timestamp - self._last_sell_timestamp) <= self._min_fire_period

  def _within_buy_update_period(self, timestamp):
    return (timestamp - self._last_buy_timestamp) <= self._min_fire_period

  '''Passive'''

  def post(self,
           timestamp,
           sell_price,
           buy_price,
           sell_qty=None,
           buy_qty=None,
           sell_price_pull=None,
           buy_price_pull=None):
    self.post_tspq = []
    self.can_tspq = []
    fired = False

    sell_qty = sell_qty if sell_qty is not None else self._get_lot_size(timestamp)
    buy_qty = buy_qty if buy_qty is not None else self._get_lot_size(timestamp)

    sell_cancel = False
    if self._ask_matcher.qty > 0:
      if (sell_price_pull is not None and self._ask_matcher.price <= sell_price_pull):
        sell_cancel = True
      elif (self._ask_matcher.price != sell_price
            and not self._within_sell_update_period(timestamp)):
        sell_cancel = True
      else:
        pass
    if not self._ask_matcher.dead and sell_cancel:
      cancel_ret = self._ask_matcher.cancel(timestamp)
      if cancel_ret:
        self._last_sell_timestamp = timestamp
        self.can_tspq.append(cancel_ret)
      fired = True

    buy_cancel = False
    if self._bid_matcher.qty > 0:
      if (buy_price_pull is not None and self._bid_matcher.price >= buy_price_pull):
        buy_cancel = True
      elif (self._bid_matcher.price != buy_price and not self._within_buy_update_period(timestamp)):
        buy_cancel = True
      else:
        pass
    if not self._bid_matcher.dead and buy_cancel:
      cancel_ret = self._bid_matcher.cancel(timestamp)
      if cancel_ret:
        self._last_buy_timestamp = timestamp
        self.can_tspq.append(cancel_ret)
      fired = True

    no_sell = False
    if sell_cancel:
      no_sell = True
    elif self._within_sell_update_period(timestamp):
      no_sell = True
    elif (sell_qty == 0 or sell_price is None
          or (self._get_min_pos(timestamp) is not None
              and self.pos.pos <= self._get_min_pos(timestamp))):
      no_sell = True
    else:
      pass

    no_buy = False
    if buy_cancel:
      no_buy = True
    elif self._within_buy_update_period(timestamp):
      no_buy = True
    elif (buy_qty == 0 or buy_price is None or (self._get_max_pos(timestamp) is not None
                                                and self.pos.pos >= self._get_max_pos(timestamp))):
      no_buy = True
    else:
      pass

    if not no_sell:
      self._last_sell_timestamp = timestamp
      self._ask_matcher.set_order(sell_price, sell_qty, timestamp)
      self.pos.submit(timestamp=timestamp, price=sell_price, qty=sell_qty, side=OrderSide.SELL)
      self.post_tspq.append((timestamp, OrderSide.SELL, sell_price, sell_qty))
      fired = True

    if not no_buy:
      self._last_buy_timestamp = timestamp
      self._bid_matcher.set_order(buy_price, buy_qty, timestamp)
      self.pos.submit(timestamp=timestamp, price=buy_price, qty=buy_qty, side=OrderSide.BUY)
      self.post_tspq.append((timestamp, OrderSide.BUY, buy_price, buy_qty))
      fired = True

    if fired:
      self._last_fire_time = timestamp
    return fired

  def update_book(self, book):
    if not (book.has_ask() and book.has_bid()):
      return

    if self._agg_executor:
      self._agg_executor.update_book(book)

    ask_fills = self._ask_matcher.update_book(book)
    bid_fills = self._bid_matcher.update_book(book)
    self._on_fill(book.timestamp, *(ask_fills + bid_fills))
    self.pos.update_book(book)

  def update_trade(self, trade):
    if self._agg_executor:
      self._agg_executor.update_trade(trade)

    trade_qty = self._trade_qty_func(trade)
    ask_fills = self._ask_matcher.update_trade(trade.price, trade_qty, trade.side, trade.timestamp)
    bid_fills = self._bid_matcher.update_trade(trade.price, trade_qty, trade.side, trade.timestamp)
    self._on_fill(trade.timestamp, *(ask_fills + bid_fills))

  def _on_fill(self, timestamp, *fills):
    for fill in fills:
      if fill is not None:
        self.fill_callback(fill.price, fill.qty, fill.side, fill.maker, timestamp)
        self.pos.update_by_fill_object(fill)
        self.pos_callback(self.pos, fill.price, fill.qty, fill.side, fill.maker, timestamp)

  def get_summary(self, postfix=None):
    return self.pos.get_summary(postfix)


class AggressionSimExecutor(object):
  def __init__(self,
               maker_fee,
               taker_fee,
               execution_delay,
               max_qty,
               exclude_on_the_fly_qty=True,
               trade_qty_func=None,
               unavail_checker=None):
    self._trade_qty_func = trade_qty_func or (lambda trade: trade.qty)
    self._execution_delay = execution_delay

    assert max_qty > 0
    self._max_qty = max_qty

    self._exclude_on_the_fly_qty = exclude_on_the_fly_qty
    self._unavail_checker = unavail_checker

    self.fill_callback = lambda price, qty, side, maker, timestamp: None
    self.pos_callback = lambda pos, price, qty, side, maker, timestamp: None

    self.pos = PositionCalculator(maker_fee, taker_fee)
    self._matchers = []

  def submit(self, timestamp, price, qty, side):
    assert qty > 0
    if self._exclude_on_the_fly_qty:
      on_the_fly_qty_sum = 0
      for matcher in self._matchers:
        if matcher.on_the_fly and matcher.side == side:
          on_the_fly_qty_sum += matcher.on_the_fly_qty
      qty -= on_the_fly_qty_sum
      if qty <= 0:
        return

    qty = min(qty, self._max_qty)
    matcher = SimpleMatcher(side,
                            self._execution_delay,
                            post_only=False,
                            unavail_checker=self._unavail_checker)
    matcher.set_order(price, qty, timestamp)
    self.pos.submit(timestamp=timestamp, price=price, qty=qty, side=side)
    self._matchers.append(matcher)

  def sell(self, timestamp, price, qty):
    self.submit(timestamp, price, qty, OrderSide.SELL)

  def buy(self, timestamp, price, qty):
    self.submit(timestamp, price, qty, OrderSide.BUY)

  def has_on_the_fly_orders(self):
    for matcher in self._matchers:
      if matcher.on_the_fly:
        return True
    return False

  def cancel_all(self, timestamp):
    for matcher in self._matchers:
      if not matcher.dead and not matcher.on_the_fly:
        matcher.cancel(timestamp)

  def cancel_condition(self, timestamp, side, minp, maxp):
    for matcher in self._matchers:
      if not matcher.dead and not matcher.on_the_fly:
        if (matcher.side == side and (minp is None or minp <= matcher.price)
            and (maxp is None or matcher.price <= maxp)):
          matcher.cancel(timestamp)

  def update_book(self, book):
    if not (book.has_ask() and book.has_bid()):
      return

    fills = []
    for matcher in self._matchers:
      if not matcher.dead:
        fills += matcher.update_book(book)
    self._on_fill(book.timestamp, *fills)
    self._clean_dead_orders()
    self.pos.update_book(book)

  def update_trade(self, trade):
    trade_qty = self._trade_qty_func(trade)
    fills = []
    for matcher in self._matchers:
      if not matcher.dead:
        fills += matcher.update_trade(trade.price, trade_qty, trade.side, trade.timestamp)
    self._on_fill(trade.timestamp, *fills)
    self._clean_dead_orders()

  def _on_fill(self, timestamp, *fills):
    for fill in fills:
      if fill is not None:
        self.fill_callback(fill.price, fill.qty, fill.side, fill.maker, timestamp)
        self.pos.update_by_fill_object(fill)
        self.pos_callback(self.pos, fill.price, fill.qty, fill.side, fill.maker, timestamp)

  def _clean_dead_orders(self):
    if len(self._matchers) < 20:
      return
    matchers = []
    for matcher in self._matchers:
      if not matcher.dead:
        matchers.append(matcher)
    self._matchers = matchers

  def get_summary(self, postfix=None):
    return self.pos.get_summary(postfix)


class PassAggSimExecutor(object):
  def __init__(self,
               pass_product,
               agg_product,
               net_pos_allowance,
               passive_param,
               aggression_param,
               pos_callback=None):
    self._products = {}
    self._executors = {}
    self._books = {}

    self._net_pos_allowance = net_pos_allowance
    self._agg_sell_price = None
    self._agg_buy_price = None
    self._pass_min_sell_fill_price = None
    self._pass_max_buy_fill_price = None

    self.pos_callback = (pos_callback
                         or (lambda exchange_id, pos, price, qty, side, maker, timestamp: None))

    self._init_passive_executor(pass_product, **passive_param)
    self._init_aggression_executor(agg_product, **aggression_param)

  def _init_passive_executor(self,
                             product,
                             lot_size,
                             min_pos,
                             max_pos,
                             maker_fee,
                             taker_fee,
                             execution_delay_sec=1.0,
                             post_only=True,
                             trade_qty_func=None):
    exchange_id = 2  # Passive
    self._products[exchange_id] = product
    self._books[exchange_id] = None

    pass_executor = PassiveSimExecutor(lot_size,
                                       min_pos,
                                       max_pos,
                                       maker_fee,
                                       taker_fee,
                                       int(execution_delay_sec * (10**9)),
                                       post_only,
                                       trade_qty_func=trade_qty_func)
    pass_executor.fill_callback = self._on_pass_fill
    pass_executor.pos_callback = functools.partial(self.pos_callback, exchange_id)
    self._executors[exchange_id] = pass_executor

  def _init_aggression_executor(self,
                                product,
                                maker_fee,
                                taker_fee,
                                max_qty,
                                execution_delay_sec=1.0,
                                trade_qty_func=None):
    exchange_id = 1  # Aggression
    self._products[exchange_id] = product
    self._books[exchange_id] = None

    agg_executor = AggressionSimExecutor(maker_fee,
                                         taker_fee,
                                         int(execution_delay_sec * (10**9)),
                                         max_qty,
                                         trade_qty_func)
    agg_executor.pos_callback = functools.partial(self.pos_callback, exchange_id)
    self._executors[exchange_id] = agg_executor

  @property
  def _pass_executor(self):
    return self._executors[2]

  @property
  def _agg_executor(self):
    return self._executors[1]

  def get_pos(self, exchange_id):
    return self._executors[exchange_id].pos

  @property
  def pnl_sum(self):
    return self.get_pos(1).pnl + self.get_pos(2).pnl

  @property
  def net_pos(self):
    return self.get_pos(1).pos + self.get_pos(2).pos

  def get_summary(self):
    return {
        'pnl': self.pnl_sum,
        'qty_sum': self.get_pos(1).qty_sum + self.get_pos(2).qty_sum,
        **self.get_pos(1).get_summary('1'),
        **self.get_pos(2).get_summary('2')
    }

  def update_book(self, exchange_id, book):
    self._books[exchange_id] = book
    self._executors[exchange_id].update_book(book)

  def update_trade(self, exchange_id, trade):
    self._executors[exchange_id].update_trade(trade)

  def _on_pass_fill(self, price, qty, side, maker, timestamp):
    assert qty > 0, qty

    agg_timestamp = timestamp + (10**6)  # Processing time: 1ms
    # max_qty is considered in the aggression executor
    if side == OrderSide.SELL:
      pos = self.net_pos - qty
      if (self._agg_buy_price is None or pos >= 0
          or price < self._pass_min_sell_fill_price):  # No hedge.
        return
      self._agg_executor.buy(agg_timestamp, self._agg_buy_price, -pos)

    elif side == OrderSide.BUY:
      pos = self.net_pos + qty
      if (self._agg_sell_price is None or pos <= 0
          or price > self._pass_max_buy_fill_price):  # No hedge.
        return
      self._agg_executor.sell(agg_timestamp, self._agg_sell_price, pos)

    else:
      raise ValueError(side)

  def set_execution(self,
                    timestamp,
                    pass_sell_price,
                    pass_buy_price,
                    agg_sell_price,
                    agg_buy_price,
                    pass_min_sell_fill_price,
                    pass_max_buy_fill_price):
    self._agg_sell_price = agg_sell_price
    self._agg_buy_price = agg_buy_price
    self._pass_min_sell_fill_price = pass_min_sell_fill_price
    self._pass_max_buy_fill_price = pass_max_buy_fill_price

    pass_sell_qty = None
    if (self._net_pos_allowance is not None and self.net_pos <= -self._net_pos_allowance):
      pass_sell_qty = 0

    pass_buy_qty = None
    if (self._net_pos_allowance is not None and self.net_pos >= self._net_pos_allowance):
      pass_buy_qty = 0

    if self._pass_executor.post(timestamp,
                                pass_sell_price,
                                pass_buy_price,
                                sell_qty=pass_sell_qty,
                                buy_qty=pass_buy_qty):
      self._agg_executor.cancel_all(timestamp)
